import os

from flask import Flask, request, jsonify
from langchain.schema import HumanMessage
from langchain_openai import ChatOpenAI

app = Flask(__name__)

# ---------- 关键配置 ----------
# SiliconFlow硅基流动 兼容 OpenAI 接口
os.environ["OPENAI_API_KEY"] = "sk-rnqurwrjarpxjetwkcphjymkcrbxpoyhpueaxwtkevmdwhlw"
os.environ["OPENAI_API_BASE"] = "https://api.siliconflow.cn/v1"

MODEL_NAME = "Qwen/QwQ-32B"

# LangChain 客户端
llm = ChatOpenAI(
    model_name=MODEL_NAME,
    temperature=0.7,
    max_tokens=512
)

# ---------- 内存历史 ----------
user_histories: dict[str, list] = {}  # list[HumanMessage | AIMessage]


def get_history_list(user_id: str):
    """返回 LangChain 消息对象列表"""
    return user_histories.setdefault(user_id, [])


@app.route("/chat", methods=["POST"])
def chat():
    data = request.get_json()
    user_id = data.get("user_id")
    user_text = data.get("message")
    if not user_id or not user_text:
        return jsonify({"error": "缺少 user_id 或 message"}), 400

    history = get_history_list(user_id)
    history.append(HumanMessage(content=user_text))

    # 调用模型
    ai_msg = llm.invoke(history)  # 返回 AIMessage
    history.append(ai_msg)

    # 转成普通 dict 给前端
    serializable = [
        {"role": "user" if isinstance(msg, HumanMessage) else "assistant",
         "content": msg.content}
        for msg in history
    ]
    return jsonify({
        "user_id": user_id,
        "response": ai_msg.content,
        "chat_history": serializable
    })


@app.route("/history/<user_id>", methods=["GET"])
def get_history(user_id):
    history = get_history_list(user_id)
    serializable = [
        {"role": "user" if isinstance(msg, HumanMessage) else "assistant",
         "content": msg.content}
        for msg in history
    ]
    return jsonify(serializable)


@app.route("/history/<user_id>", methods=["DELETE"])
def clear_history(user_id):
    user_histories.pop(user_id, None)
    return jsonify({"message": f"用户 {user_id} 的聊天历史已清除"})


@app.route("/")
def index():
    return jsonify({"message": "Flask + LangChain 聊天服务已启动",
                    "endpoints": ["/chat", "/history/<user_id>"]})


if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=8000)
